import pybullet as p
import pybullet_data
import time
import pygame


user_control_force = 15

p.connect(p.GUI)
p.setGravity(0, 0, -10)


# 设置相机参数
camera_distance = 2   # 相机到目标的距离
camera_yaw = 45          # 水平旋转角度（偏航角）
camera_pitch = -30       # 垂直旋转角度（俯仰角）
camera_target = [0, 0, 1.05] # 相机对准的目标位置

# 重置相机视角
p.resetDebugVisualizerCamera(
    cameraDistance=camera_distance,
    cameraYaw=camera_yaw,
    cameraPitch=camera_pitch,
    cameraTargetPosition=camera_target
)
# 创建地面
p.createCollisionShape(p.GEOM_PLANE)
p.createMultiBody(0, 0)

# 加载机械臂
robot_id = p.loadURDF("./urdf/simple_arm.urdf", [0, 0, 0.001])

# 获取关节信息
num_joints = p.getNumJoints(robot_id)
print(f"机械臂有 {num_joints} 个关节")



# 控制第一个关节
p.setJointMotorControl2(
    bodyUniqueId=robot_id,
    jointIndex=0,
    controlMode=p.VELOCITY_CONTROL,
    force=0
)

# 控制第二个关节
p.setJointMotorControl2(
    bodyUniqueId=robot_id,
    jointIndex=1,
    controlMode=p.VELOCITY_CONTROL,
    force=0
)



for i in p.getDebugVisualizerCamera():
    print(i)
# p.createConstraint(parentBodyUniqueId=robot_id, parentLinkIndex=0, -1, -1, jointType=p.JOINT_FIXED, [0,0,0],[0,0,0],[0,0,0])
p.createConstraint(parentBodyUniqueId=robot_id, parentLinkIndex=-1, childBodyUniqueId=-1,childLinkIndex=-1,jointType=p.JOINT_FIXED, jointAxis=[0,0,0], parentFramePosition=[0,0,-0.0005], childFramePosition=camera_target)


KEY_A = ord("a")
KEY_D = ord("d")
KEY_Q = ord("q")
# 仿真循环
try:
    while True:
        
        events = p.getKeyboardEvents()
        # PyBullet 里按下是 p.KEY_IS_DOWN，抬起是 p.KEY_WAS_RELEASED
        a_down = events.get(KEY_A, 0) & p.KEY_IS_DOWN
        d_down = events.get(KEY_D, 0) & p.KEY_IS_DOWN
        esc_down = events.get(KEY_Q, 0) & p.KEY_IS_DOWN
        if a_down:
            p.setJointMotorControl2(robot_id, 0, p.TORQUE_CONTROL, force=user_control_force)
        elif d_down:
            p.setJointMotorControl2(robot_id, 0, p.TORQUE_CONTROL, force=-user_control_force)

        elif esc_down:
            # 也可以用 KEY_IS_DOWN：持续按住为真；WAS_TRIGGERED：按下瞬间为真
            break
        else:
            p.setJointMotorControl2(robot_id, 0, p.TORQUE_CONTROL, force=0.0)
            
        p.stepSimulation()
        time.sleep(1./240.)
finally:
    p.disconnect()